# adk_mcp_server.py
import asyncio
import json
from dotenv import load_dotenv

# MCP 服务器导入
from mcp import types as mcp_types # 使用别名以避免与 genai.types 冲突
from mcp.server.lowlevel import Server, NotificationOptions
from mcp.server.models import InitializationOptions
import mcp.server.stdio

# ADK 工具导入
from google.adk.tools.function_tool import FunctionTool
from google.adk.tools.load_web_page import load_web_page # 示例 ADK 工具
# ADK <-> MCP 转换工具
from google.adk.tools.mcp_tool.conversion_utils import adk_to_mcp_tool_type

# --- 加载环境变量（如果 ADK 工具需要它们） ---
load_dotenv()

# --- 准备 ADK 工具 ---
# 实例化你想要公开的 ADK 工具
print("初始化 ADK load_web_page 工具...")
adk_web_tool = FunctionTool(load_web_page)
print(f"ADK 工具 '{adk_web_tool.name}' 已初始化。")
# --- ADK 工具准备结束 ---

# --- MCP 服务器设置 ---
print("创建 MCP 服务器实例...")
# 创建一个命名的 MCP 服务器实例
app = Server("adk-web-tool-mcp-server")

# 实现 MCP 服务器的 @app.list_tools 处理程序
@app.list_tools()
async def list_tools() -> list[mcp_types.Tool]:
  """MCP 处理程序，列出可用工具。"""
  print("MCP 服务器：收到 list_tools 请求。")
  # 将 ADK 工具的定义转换为 MCP 格式
  mcp_tool_schema = adk_to_mcp_tool_type(adk_web_tool)
  print(f"MCP 服务器：公布工具：{mcp_tool_schema.name}")
  return [mcp_tool_schema]

# 实现 MCP 服务器的 @app.call_tool 处理程序
@app.call_tool()
async def call_tool(
    name: str, arguments: dict
) -> list[mcp_types.TextContent | mcp_types.ImageContent | mcp_types.EmbeddedResource]:
  """MCP 处理程序，执行工具调用。"""
  print(f"MCP 服务器：收到 call_tool 请求，用于 '{name}'，参数：{arguments}")

  # 检查请求的工具名称是否与我们包装的 ADK 工具匹配
  if name == adk_web_tool.name:
    try:
      # 执行 ADK 工具的 run_async 方法
      # 注意：tool_context 为 None，因为我们不在完整的 ADK Runner 调用中
      adk_response = await adk_web_tool.run_async(
          args=arguments,
          tool_context=None, # 这里没有可用的 ADK 上下文
      )
      print(f"MCP 服务器：ADK 工具 '{name}' 成功执行。")
      # 将 ADK 工具的响应（通常是字典）格式化为 MCP 格式。
      # 这里，我们在 TextContent 中将响应字典序列化为 JSON 字符串。
      # 根据特定 ADK 工具的输出和客户端需求调整格式化。
      response_text = json.dumps(adk_response, indent=2)
      return [mcp_types.TextContent(type="text", text=response_text)]

    except Exception as e:
      print(f"MCP 服务器：执行 ADK 工具 '{name}' 时出错：{e}")
      # 以 MCP 格式返回错误消息
      # 创建适当的 MCP 错误响应可能更为健壮
      error_text = json.dumps({"error": f"无法执行工具 '{name}'：{str(e)}"})
      return [mcp_types.TextContent(type="text", text=error_text)]
  else:
      # 处理对未知工具的调用
      print(f"MCP 服务器：未找到工具 '{name}'。")
      error_text = json.dumps({"error": f"未实现工具 '{name}'。"})
      # 为简单起见，将错误作为 TextContent 返回
      return [mcp_types.TextContent(type="text", text=error_text)]

# --- MCP 服务器运行器 ---
async def run_server():
  """通过标准输入/输出运行 MCP 服务器。"""
  # 使用 MCP 库中的 stdio_server 上下文管理器
  async with mcp.server.stdio.stdio_server() as (read_stream, write_stream):
    print("MCP 服务器开始握手...")
    await app.run(
        read_stream,
        write_stream,
        InitializationOptions(
            server_name=app.name, # 使用上面定义的服务器名称
            server_version="0.1.0",
            capabilities=app.get_capabilities(
                # 定义服务器功能 - 查阅 MCP 文档了解选项
                notification_options=NotificationOptions(),
                experimental_capabilities={},
            ),
        ),
    )
    print("MCP 服务器运行循环已完成。")

if __name__ == "__main__":
  print("启动公开 ADK 工具的 MCP 服务器...")
  try:
    asyncio.run(run_server())
  except KeyboardInterrupt:
    print("\nMCP 服务器被用户停止。")
  except Exception as e:
    print(f"MCP 服务器遇到错误：{e}")
  finally:
    print("MCP 服务器进程退出。")
# --- MCP 服务器结束 ---